# -*- coding: utf-8 -*-
'''
sklearn实现贝叶斯文本分类
'''

import pandas as pd
from sklearn.preprocessing import LabelEncoder
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.model_selection import train_test_split
from sklearn.externals import joblib
from sklearn.naive_bayes import MultinomialNB
from sklearn import metrics

df = pd.read_csv('.//data//SMSSpamCollection',delimiter='\t',header=None)
tfidf_vec = TfidfVectorizer(stop_words='english')

X = df[1].values
y = df[0].values

label_encoder = LabelEncoder()
y = label_encoder.fit_transform(y)

tfidf_vec.fit(X)
joblib.dump(tfidf_vec,'tfidf_vec.model')

X = tfidf_vec.transform(X)
X_train, X_test, y_train, y_test = train_test_split(X,y)

# 训练模型
clf = MultinomialNB()
clf.fit(X_train, y_train)
joblib.dump(clf, 'clf.model')

# 测试集表现
predicted = clf.predict(X_test)
print(metrics.classification_report(y_test, predicted,
    target_names=['ham','spam']))


# 模型复用
tfidf_vec = joblib.load('tfidf_vec.model')
clf = joblib.load('clf.model')
test = 'Go until jurong point, crazy.. Available only in bugis n great world la e buffet... Cine there got amore wat...'
predicted = clf.predict(tfidf_vec.transform([test]))














